import { slice, mul, concat2d } from "@tensorflow/tfjs-core";
import type { Tensor1D, Tensor2D } from "@tensorflow/tfjs-core";

export type Box = {
  startEndTensor: Tensor2D;
  startPoint: Tensor2D;
  endPoint: Tensor2D;
};

export const disposeBox = (box: Box): void => {
  box.startEndTensor.dispose();
  box.startPoint.dispose();
  box.endPoint.dispose();
};

export const createBox = (startEndTensor: Tensor2D): Box => ({
  startEndTensor,
  startPoint: slice(startEndTensor, [0, 0], [-1, 2]),
  endPoint: slice(startEndTensor, [0, 2], [-1, 2]),
});

export const scaleBox = (box: Box, factors: Tensor1D | [number, number]) => {
  const starts = mul(box.startPoint, factors);
  const ends = mul(box.endPoint, factors);
  const newCoordinates = concat2d([starts as Tensor2D, ends as Tensor2D], 1);

  return createBox(newCoordinates);
};
